from nltk.corpus import verbnet
import themrole
import framesyntax

class verb(object):
   def __init__(self, plainword):
      self.value = plainword

      #not sure about this, but it's a way to get a frame for 'to do'
      if plainword in {'do', 'does', 'did', 'did\'t', 'doesn\'t'}:
         plainword = 'begin'

      self.senses = \
         [verbsense(s) for s in verbnet.classids(plainword)]

      '''extend the explicitly provided senses/classes with their ancestors'''
      for sense in self.senses:
         print sense
         hyphenscount = sense.classid.count('-')
         if hyphenscount > 1:  # two hypens implies subclass
            hyphenindex = sense.classid.rindex('-')
            superclassid = sense.classid[0:hyphenindex]
            print superclassid
            self.senses.append(verbsense(superclassid))

   '''return all frames from all senses which contain a particular prep'''
   def frameswithprep(self, prepstring):
      frames = []
      for sense in self.senses:
         frames.extend(sense.frameswithprep(prepstring))
      return frames

   def __repr__(self):
      return "%s('%s')" % (type(self).__name__, self.value)

class verbsense(object):
   def __init__(self, vnclassid):
      '''
      Constructor
      Takes in a verbnet classid
      '''
      vntree = verbnet.vnclass(vnclassid)
      self.themroles = dict()
      self.syntaxes = list()
      self.classid = vntree.get('ID')

      '''if this is a subclass, APPARENTLY I have to get the superclass
         thematic role selectional restrictions MANUALLY.'''
      if vntree.get('ID').count('-') > 1:  # two hypens implies subclass
         # grab the part before the second hypen
         superclassname = vntree.get('ID')[0:vntree.get('ID').rindex('-')]
         superclass = verbnet.vnclass(superclassname)
         '''from the superclass, get thematic roles
            and their related restrictions'''
         roles = superclass.find("THEMROLES")
         for role in roles:
            newrole = themrole.themrole(role)
            self.themroles[newrole.roletype] = newrole

      '''get the thematic roles and their related restrictions'''
      roles = vntree.find("THEMROLES")
      for role in roles:
         newrole = themrole.themrole(role)
         self.themroles[newrole.roletype] = newrole
         '''parent roles are overwritten with new restrictions'''
      for tr in self.themroles:
         print tr
         print "Require: ", self.themroles[tr].require
         print "Refuse: ", self.themroles[tr].refuse

      '''get the frame syntaxes'''
      frames = vntree.find("FRAMES")
      for frame in frames:
         syntax = framesyntax.framesyntax(frame.find("SYNTAX"), self.themroles)
         self.syntaxes.append(syntax)

   '''return the frames from this sense which contain a particular prep'''
   def frameswithprep(self, prepstring):
      frames = []
      for frame in self.syntaxes:
         if frame.hasprep(prepstring):
            frames.append(frame)
      return frames

   def __repr__(self):
      return "%s('%s')" % (type(self).__name__, self.classid)

